#include <arrow/api.h>
#include <arrow/io/api.h>
#include <parquet/arrow/reader.h>
#include <iostream>

int main() {
    // 1️⃣ 打开 Parquet 文件
    auto infile_result = arrow::io::ReadableFile::Open("example.parquet");
    if (!infile_result.ok()) {
        std::cerr << "❌ 无法打开输入文件: " << infile_result.status().ToString() << std::endl;
        return -1;
    }
    std::shared_ptr<arrow::io::ReadableFile> infile = infile_result.ValueOrDie();

    // 2️⃣ 创建 Parquet Arrow 文件读取器（新版 API ✅）
    auto reader_result = parquet::arrow::OpenFile(infile, arrow::default_memory_pool());
    if (!reader_result.ok()) {
        std::cerr << "❌ 无法创建 Parquet 读取器: " << reader_result.status().ToString() << std::endl;
        return -1;
    }
    std::unique_ptr<parquet::arrow::FileReader> reader = std::move(reader_result.ValueOrDie());

    // 3️⃣ 读取整个文件到 Arrow Table
    std::shared_ptr<arrow::Table> table;
    auto status = reader->ReadTable(&table);
    if (!status.ok()) {
        std::cerr << "❌ 读取 Parquet 文件失败: " << status.ToString() << std::endl;
        return -1;
    }

    std::cout << "✅ 成功读取 Parquet 文件: example.parquet" << std::endl;
    std::cout << "表的列数: " << table->num_columns() << std::endl;
    std::cout << "表的行数: " << table->num_rows() << std::endl;

    // 4️⃣ 打印表数据
    for (int c = 0; c < table->num_columns(); ++c) {
        auto column = table->column(c);
        std::cout << "列 " << c << " 名称: " << table->schema()->field(c)->name() << std::endl;

        for (int chunk = 0; chunk < column->num_chunks(); ++chunk) {
            auto array = column->chunk(chunk);

            if (array->type_id() == arrow::Type::INT64) {
                auto int_array = std::static_pointer_cast<arrow::Int64Array>(array);
                for (int64_t i = 0; i < int_array->length(); ++i) {
                    if (!int_array->IsNull(i)) std::cout << int_array->Value(i) << " ";
                    else std::cout << "NULL ";
                }
            } else if (array->type_id() == arrow::Type::STRING) {
                auto str_array = std::static_pointer_cast<arrow::StringArray>(array);
                for (int64_t i = 0; i < str_array->length(); ++i) {
                    if (!str_array->IsNull(i)) std::cout << str_array->GetString(i) << " ";
                    else std::cout << "NULL ";
                }
            }
            std::cout << std::endl;
        }
    }

    return 0;
}
